#include "oppc.hpp"

namespace oppc_internal { // 辅助函数和内部模块仍然在此命名空间内

// --- 调试辅助函数：打印 ap_fixed 值而不使用 to_float() ---
#ifndef __SYNTHESIS__
// 辅助函数，用于打印 ap_fixed 的原始位模式和其十进制表示 (通过 operator<<)
template<int W, int I, ap_q_mode Q, ap_o_mode O, int N>
void print_fixed_value(const std::string& name, const ap_fixed<W, I, Q, O, N>& val) {
    std::cout << name << " (W=" << W << ", I=" << I << ", F=" << (W - I) << "): "
              << val << " (Raw bits: 0x" << std::hex << val.to_uint() << std::dec << ")" << std::endl;
}

// 辅助函数，用于获取 ap_fixed 类型的 LSB 值
template<typename T>
T get_lsb_fixed_value() {
    T lsb_val = 0; // 所有位清零
    lsb_val.range(0, 0) = 1; // 将最低位设置为 1
    return lsb_val;
}
#endif
// --- 调试辅助函数结束 ---

// -----------------------------------------------------------
// 纹理缓存管理模块 (点采样版本)
// -----------------------------------------------------------

// 纹理缓存块加载函数
void load_texture_block(
    TextureCacheLine cache[TEXTURE_CACHE_LINES], // 接收片上缓存指针
    ap_uint<32>* texture_axi,                               // DDR 纹理数据接口
    ap_uint<18> block_addr_au,                                 // 要加载的纹理块在DDR中的逻辑地址，修正位宽
    ap_uint<7> cache_line_idx_au                              // 缓存行索引
) {
    #pragma HLS INLINE off
    #pragma HLS PIPELINE II=1 // 内部循环展开，这个函数整体可以流水线化

    ap_uint<32> ddr_start_offset_calc = (ap_uint<32>)block_addr_au * TEX_CACHE_BURST_LEN; // 使用乘法，但这里是整数乘法，不需要 fpal::mul
    unsigned int cache_line_idx = cache_line_idx_au.to_uint();
    
    // 从DDR突发读取数据
    for (int i = 0; i < TEX_CACHE_BURST_LEN; ++i) { 
        #pragma HLS PIPELINE II=1
        cache[cache_line_idx].data[i] = texture_axi[ddr_start_offset_calc + i]; 
    }

    // 更新缓存行状态
    cache[cache_line_idx].valid_bit = true;
    cache[cache_line_idx].tag = block_addr_au;
}

// 计算缓存行索引 (直接映射)
ap_uint<7> get_cache_line_idx(ap_uint<18> block_addr) { // 修正参数位宽
    #pragma HLS INLINE
    return block_addr % TEXTURE_CACHE_LINES;
}

// 获取单个纹素值（包含缓存命中/不命中逻辑）- 点采样版本
ap_uint<32> get_texel_value_point_sample(
    TextureCacheLine cache[TEXTURE_CACHE_LINES],
    ap_uint<32>* texture_axi,
    ap_uint<16> u, ap_uint<16> v
) {
    #pragma HLS INLINE // 内联此函数，HLS 能更好地优化缓存逻辑
    
    // global_tex_idx_au 的最大值 (2047 * 2048 + 2047) 约为 4*10^6，需要 23 位
    ap_uint<23> global_tex_idx_au = (ap_uint<23>)v * TEX_WIDTH + u; // 整数乘法/加法，不需要 fpal::
    
    // block_addr_au 的最大值 (4194303 / 16) 约为 262143，需要 18 位
    ap_uint<18> block_addr_au = global_tex_idx_au / TEX_CACHE_BURST_LEN; // 整数除法，不需要 fpal::
    ap_uint<4> block_offset_au = global_tex_idx_au % TEX_CACHE_BURST_LEN; // 整数模运算，不需要 fpal::
    
    ap_uint<7> cache_line_idx_au = get_cache_line_idx(block_addr_au); // 直接映射缓存行索引
    
    unsigned int cache_line_idx = cache_line_idx_au.to_uint();
    unsigned int block_offset = block_offset_au.to_uint();

    // 缓存命中检测
    if (!cache[cache_line_idx].valid_bit || cache[cache_line_idx].tag != block_addr_au) {
        // 缓存不命中，从DDR加载整个块
        load_texture_block(cache, texture_axi, block_addr_au, cache_line_idx_au);
    }
    return cache[cache_line_idx].data[block_offset];
}


// 纹理采样函数 (点采样)
ap_uint<24> sample_texture_point_sample(
    TextureCacheLine cache[TEXTURE_CACHE_LINES], // 接收片上缓存指针
    ap_uint<32>* texture_axi,                               // DDR 纹理数据接口
    fpal::fp_vec_unit_t uv_u,
    fpal::fp_vec_unit_t uv_v
) {
    #pragma HLS INLINE // 将此函数内联到 process_single_pixel，简化 DATAFLOW 调度，减少函数调用开销

    // UV 坐标映射到纹素坐标范围 [0, TEXTURE_SIZE-1]
    fpal::fp_vec_unit_t clamped_uv_u = fpal::clamp(uv_u, (fpal::fp_vec_unit_t)0.0, (fpal::fp_vec_unit_t)1.0);
    fpal::fp_vec_unit_t clamped_uv_v = fpal::clamp(uv_v, (fpal::fp_vec_unit_t)0.0, (fpal::fp_vec_unit_t)1.0);

    // 将 UV 转换为整数纹素坐标。对于点采样，直接取整即可。
    // 这里使用 fp_coord_t 进行中间计算，因为它有更大的整数部分范围，
    // 以防止 (TEX_WIDTH-1) 或 (TEX_HEIGHT-1) 乘以 fp_vec_unit_t 导致溢出。
    // fp_coord_t (I=14) 可容纳 2047。
    fpal::fp_coord_t tex_x_fp = fpal::mul(fpal::conv_vec_unit_to_coord(clamped_uv_u), (fpal::fp_coord_t)(TEX_WIDTH - 1));
    fpal::fp_coord_t tex_y_fp = fpal::mul(fpal::conv_vec_unit_to_coord(clamped_uv_v), (fpal::fp_coord_t)(TEX_HEIGHT - 1));

    unsigned int u_int_raw = tex_x_fp.to_uint();
    unsigned int v_int_raw = tex_y_fp.to_uint();

    // 钳制到有效纹理坐标范围 [0, TEX_SIZE-1]
    ap_uint<16> u_texel = (ap_uint<16>)fpal::clamp_int(u_int_raw, (unsigned int)0, (unsigned int)(TEX_WIDTH - 1));
    ap_uint<16> v_texel = (ap_uint<16>)fpal::clamp_int(v_int_raw, (unsigned int)0, (unsigned int)(TEX_HEIGHT - 1));
    
    // 获取单个纹素值
    ap_uint<32> texel_value_32bit = get_texel_value_point_sample(cache, texture_axi, u_texel, v_texel);

    // 提取 24 位 RGB 颜色
    ap_uint<24> final_color;
    final_color.range(23, 16) = texel_value_32bit.range(23, 16); // Red
    final_color.range(15, 8) = texel_value_32bit.range(15, 8);   // Green
    final_color.range(7, 0) = texel_value_32bit.range(7, 0);     // Blue
    
    return final_color;
}


// Blinn-Phong 光照模型计算
ap_uint<24> calculate_blinn_phong(
    const OppeConfig& config,
    const fpal::vec3_fp_vec_unit& interpolated_normal,
    const ap_uint<24> texel_color_24bit
) {
    #pragma HLS INLINE // 确保内联到 process_single_pixel，减少 DATAFLOW 调用开销，使得 II=1 更容易实现
    
    // 光源方向和法线向量归一化
    fpal::vec3_fp_vec_unit light_dir_norm = fpal::vec3_normalize(config.light_dir); 
    fpal::vec3_fp_vec_unit N_norm = fpal::vec3_normalize(interpolated_normal);     

    // 假设视角方向是 (0,0,1)，即相机沿着Z轴正向看，屏幕平面是XY平面。
    // 这意味着所有像素的视角向量 V 都是 (0,0,1)。
    fpal::vec3_fp_vec_unit V = { (fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)1.0f }; 

    // 将 8 位颜色值归一化到 [0, 1] 范围的 fp_vec_unit_t
    auto conv_color_channel_to_vec_unit = [](fpal::fp_color_channel_t val) {
        return fpal::div((fpal::fp_vec_unit_t)val, (fpal::fp_vec_unit_t)255.0f);
    };

    fpal::fp_vec_unit_t light_color_r_norm = conv_color_channel_to_vec_unit(config.light_color.r);
    fpal::fp_vec_unit_t light_color_g_norm = conv_color_channel_to_vec_unit(config.light_color.g);
    fpal::fp_vec_unit_t light_color_b_norm = conv_color_channel_to_vec_unit(config.light_color.b);

    fpal::fp_vec_unit_t ambient_color_r_norm = conv_color_channel_to_vec_unit(config.ambient_color.r);
    fpal::fp_vec_unit_t ambient_color_g_norm = conv_color_channel_to_vec_unit(config.ambient_color.g);
    fpal::fp_vec_unit_t ambient_color_b_norm = conv_color_channel_to_vec_unit(config.ambient_color.b);

    fpal::fp_vec_unit_t texel_r_norm = conv_color_channel_to_vec_unit(texel_color_24bit.range(23,16));
    fpal::fp_vec_unit_t texel_g_norm = conv_color_channel_to_vec_unit(texel_color_24bit.range(15,8));
    fpal::fp_vec_unit_t texel_b_norm = conv_color_channel_to_vec_unit(texel_color_24bit.range(7,0));


    // 环境光 (Ambient)
    fpal::fp_vec_unit_t ambient_r = fpal::mul(ambient_color_r_norm, texel_r_norm);
    fpal::fp_vec_unit_t ambient_g = fpal::mul(ambient_color_g_norm, texel_g_norm);
    fpal::fp_vec_unit_t ambient_b = fpal::mul(ambient_color_b_norm, texel_b_norm);

    // 漫反射 (Diffuse)
    fpal::fp_vec_unit_t N_dot_L = fpal::vec3_dot(N_norm, light_dir_norm);
    // 钳制 N_dot_L 到 [0, 1]
    N_dot_L = fpal::clamp(N_dot_L, (fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)1.0f); 

    fpal::fp_vec_unit_t diffuse_r = fpal::mul(fpal::mul(texel_r_norm, light_color_r_norm), N_dot_L);
    fpal::fp_vec_unit_t diffuse_g = fpal::mul(fpal::mul(texel_g_norm, light_color_g_norm), N_dot_L);
    fpal::fp_vec_unit_t diffuse_b = fpal::mul(fpal::mul(texel_b_norm, light_color_b_norm), N_dot_L);

    // 镜面反射 (Specular)
    fpal::vec3_fp_vec_unit H = fpal::vec3_normalize(fpal::vec3_add(light_dir_norm, V)); // 半程向量 (Halfway vector)
    
    fpal::fp_vec_unit_t N_dot_H = fpal::vec3_dot(N_norm, H);
    // 钳制 N_dot_H 到 [0, 1]
    N_dot_H = fpal::clamp(N_dot_H, (fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)1.0f);
    
    // 增加 SHININESS 值，使高光更集中
    const int SHININESS = 12;

    // --- 使用更高精度的定点数进行 specular_intensity 的计算 ---
    // fp_specular_t 的整数位宽为 1 足够，因为 N_dot_H 已经钳制到 [0, 1]
    typedef ap_fixed<32, 1, AP_RND_CONV, AP_SAT> fp_specular_t; 
    fp_specular_t specular_intensity_high_prec = (fp_specular_t)1.0f;
    fp_specular_t N_dot_H_high_prec = (fp_specular_t)N_dot_H; // 将 N_dot_H 转换为高精度

    for (int i = 0; i < SHININESS; ++i) {
        #pragma HLS LOOP_TRIPCOUNT min=SHININESS max=SHININESS 
        #pragma HLS PIPELINE II=1
        specular_intensity_high_prec = fpal::mul(specular_intensity_high_prec, N_dot_H_high_prec); // 使用 fpal::mul 保持一致性
    }

    // --- 镜面反射强度放大因子 ---
    // 将 SPECULAR_BOOST_FACTOR 降低到更合理的值，通常是 1.0 或者一个稍大的值来补偿衰减。
    // 如果希望高光亮度更高，可以稍微增大，但不要超过 fp_vec_unit_t 的最大整数值。
    // fp_vec_unit_t 的 I=5，最大可表示大约 15.99。我们希望最终颜色饱和到 255，即 fp_vec_unit_t 饱和到 1.0。
    // 因此，boosted_specular_intensity_high_prec 最终应该在 [0, 1] 之间。
    // 所以，SPECULAR_BOOST_FACTOR 最好保持在 1.0 或接近 1.0，或者更小的数。
    // 如果 SHININESS 足够高，N_dot_H^SHININESS 就会很小，可能需要一个小的放大。
    // 这里我们尝试 2.0，给高光一个适度的亮度。
    typedef ap_fixed<16, 5, AP_RND_CONV, AP_SAT> fp_boost_factor_t; // 使用与 fp_vec_unit_t 兼容的类型
    const fp_boost_factor_t SPECULAR_BOOST_FACTOR = (fp_boost_factor_t)2.0;

    // 这里的乘法会将 specular_intensity_high_prec 转换为 fp_boost_factor_t 的类型进行运算
    // fp_specular_t (32,1) * fp_boost_factor_t (16,5) 会产生一个更宽的中间结果。
    // 结果会被截断为 fp_vec_unit_t (16,5)。
    fpal::fp_vec_unit_t specular_intensity = fpal::mul((fpal::fp_vec_unit_t)specular_intensity_high_prec, SPECULAR_BOOST_FACTOR);

    // 镜面反射颜色通常是光源颜色，不乘以纹理颜色
    fpal::fp_vec_unit_t specular_r = fpal::mul(light_color_r_norm, specular_intensity); 
    fpal::fp_vec_unit_t specular_g = fpal::mul(light_color_g_norm, specular_intensity);
    fpal::fp_vec_unit_t specular_b = fpal::mul(light_color_b_norm, specular_intensity);

    // 最终颜色 = 环境光 + 漫反射 + 镜面反射
    fpal::fp_vec_unit_t final_r_norm = fpal::add(fpal::add(ambient_r, diffuse_r), specular_r);
    fpal::fp_vec_unit_t final_g_norm = fpal::add(fpal::add(ambient_g, diffuse_g), specular_g);
    fpal::fp_vec_unit_t final_b_norm = fpal::add(fpal::add(ambient_b, diffuse_b), specular_b);

    ap_uint<24> final_color_24bit;
    final_color_24bit.range(23, 16) = fpal::conv_fp_to_color_channel(final_r_norm); // Red
    final_color_24bit.range(15, 8) = fpal::conv_fp_to_color_channel(final_g_norm);   // Green
    final_color_24bit.range(7, 0) = fpal::conv_fp_to_color_channel(final_b_norm);     // Blue

    return final_color_24bit;
}

// Tile 缓冲区初始化函数
void initialize_tile_buffers(fpal::fp_z_buffer_t z_buf[TILE_PIXELS], ap_uint<24> fb_buf[TILE_PIXELS], const int num_pixels) {
    #pragma HLS INLINE off 
    #pragma HLS ARRAY_PARTITION variable=z_buf cyclic factor=8 dim=1
    #pragma HLS ARRAY_PARTITION variable=fb_buf cyclic factor=8 dim=1
    
    for (int i = 0; i < num_pixels; ++i) {
        #pragma HLS PIPELINE II=1
        #pragma HLS LOOP_TRIPCOUNT min=TILE_PIXELS max=TILE_PIXELS
        z_buf[i] = (fpal::fp_z_buffer_t)0xFFFF; 
        fb_buf[i] = 0x000000;              
    }
}

// 纹理缓存初始化函数
void initialize_texture_cache(TextureCacheLine cache[TEXTURE_CACHE_LINES], const int cache_size) {
    #pragma HLS INLINE off
    // texture_cache_bram 已经在 oppc_core 中 complete partition，此处不需要再次指定
    // #pragma HLS ARRAY_PARTITION variable=cache complete dim=1 

    for (int i = 0; i < cache_size; ++i) {
        #pragma HLS PIPELINE II=1
        #pragma HLS LOOP_TRIPCOUNT min=TEXTURE_CACHE_LINES max=TEXTURE_CACHE_LINES

        cache[i].valid_bit = false; // 所有缓存行标记为无效
        cache[i].tag = 0;           // 标签清零
    }
}


// 处理单个像素片段的具名函数
void process_single_pixel(
    const PixelFragment& fragment,
    const OppeConfig& config,
    ap_uint<32>* texture_axi,
    ap_uint<24> frame_buffer_bram[TILE_PIXELS],
    fpal::fp_z_buffer_t z_buffer_bram[TILE_PIXELS],
    TextureCacheLine texture_cache_bram[TEXTURE_CACHE_LINES] // 传入内部静态缓存
) {
    #pragma HLS INLINE // 确保内联到 process_pixels_consumer，减少 DATAFLOW 调用开销，使得 II=1 更容易实现

    ap_uint<16> tile_x_au = fragment.screen_x - config.tile_start_x; // 整数减法
    ap_uint<16> tile_y_au = fragment.screen_y - config.tile_start_y; // 整数减法

    // 显式将边界值转换为 unsigned int
    unsigned int clamped_tile_x_val = tile_x_au.to_uint();
    unsigned int tile_x = fpal::clamp_int(clamped_tile_x_val, (unsigned int)0, (unsigned int)(TILE_WIDTH - 1));
    
    unsigned int clamped_tile_y_val = tile_y_au.to_uint();
    unsigned int tile_y = fpal::clamp_int(clamped_tile_y_val, (unsigned int)0, (unsigned int)(TILE_HEIGHT - 1));

    unsigned int pixel_idx = tile_y * TILE_WIDTH + tile_x; // 整数乘加

    fpal::fp_z_buffer_t current_z_in_buffer = z_buffer_bram[pixel_idx]; 
    
    if (fragment.interpolated_z < current_z_in_buffer) { 
        
        ap_uint<24> texel_color = oppc_internal::sample_texture_point_sample( // 调用点采样版本
            texture_cache_bram, texture_axi, fragment.interpolated_uv.u, fragment.interpolated_uv.v
        );

        ap_uint<24> final_color = oppc_internal::calculate_blinn_phong(
            config, 
            fragment.interpolated_normal, 
            texel_color
        );

        z_buffer_bram[pixel_idx] = fragment.interpolated_z;   
        frame_buffer_bram[pixel_idx] = final_color;           
    }
}

// DATAFLOW 区域的生产者函数
void read_fragments_producer(
    hls::stream<PixelFragment>& fragment_stream,
    hls::stream<PixelFragment>& valid_fragment_stream
) {
    // 为了支持 II=1，需要将 for 循环设置为固定次数。
    // 在 HLS 中，如果 stream 读取在循环条件中，可能导致 II 无法达到 1。
    // 确保生产和消费的循环次数严格一致。
    for (int i = 0; i < TILE_PIXELS; ++i) { 
        #pragma HLS PIPELINE II=1
        #pragma HLS LOOP_TRIPCOUNT min=TILE_PIXELS max=TILE_PIXELS
        PixelFragment fragment; // 声明在循环内部
        fragment_stream.read(fragment); // 阻塞读取，等待数据
        valid_fragment_stream.write(fragment); 
    }
}

// DATAFLOW 区域的消费者函数
void process_pixels_consumer(
    hls::stream<PixelFragment>& valid_fragment_stream,
    const OppeConfig& config,
    ap_uint<32>* texture_axi,
    ap_uint<24> frame_buffer_bram[TILE_PIXELS],
    fpal::fp_z_buffer_t z_buffer_bram[TILE_PIXELS],
    TextureCacheLine texture_cache_bram[TEXTURE_CACHE_LINES] // 传入内部静态缓存
) {
    for (int i = 0; i < TILE_PIXELS; ++i) { 
        #pragma HLS LOOP_TRIPCOUNT min=TILE_PIXELS max=TILE_PIXELS
        #pragma HLS PIPELINE II=1
        PixelFragment fragment; // 声明在循环内部
        valid_fragment_stream.read(fragment); // 阻塞读取
        oppc_internal::process_single_pixel(
            fragment, config, texture_axi, frame_buffer_bram, z_buffer_bram, texture_cache_bram
        );
    }
}


} // namespace oppc_internal


// -----------------------------------------------------------
// OPPC 模块的顶层函数定义
// -----------------------------------------------------------
void oppc_core(
    hls::stream<PixelFragment>& fragment_stream,
    const OppeConfig& config,
    ap_uint<32>* texture_axi, // M_AXI 纹理数据接口
    ap_uint<24> frame_buffer_bram[TILE_PIXELS], // 片上Framebuffer BRAM (作为输入/输出BRAM)
    fpal::fp_z_buffer_t z_buffer_bram[TILE_PIXELS], // 片上Z-Buffer BRAM (作为输入/输出BRAM)
    hls::stream<bool>& tile_processed_signal // 信号流，指示Tile处理完成
) {
    // 接口指令
    #pragma HLS INTERFACE s_axilite port=return bundle=AXI_LITE_OPPC_CTRL
    #pragma HLS INTERFACE axis port=fragment_stream
    #pragma HLS INTERFACE s_axilite port=config bundle=AXI_LITE_OPPC_CTRL
    // 调整 M_AXI 接口的 depth 以匹配实际纹理大小 (TEX_WIDTH * TEX_HEIGHT)
    #pragma HLS INTERFACE m_axi offset=slave port=texture_axi bundle=MAXI_TEX_DATA depth=4194304 latency=64 num_read_outstanding=32 max_read_burst_length=16

    #pragma HLS INTERFACE bram port=frame_buffer_bram
    #pragma HLS INTERFACE bram port=z_buffer_bram

    #pragma HLS ARRAY_PARTITION variable=frame_buffer_bram cyclic factor=8 dim=1
    #pragma HLS ARRAY_PARTITION variable=z_buffer_bram cyclic factor=8 dim=1

    // 纹理缓存声明为 static 局部变量，确保其在函数调用之间保持状态
    // 并进行完全分区以提供最大并行访问能力
    static TextureCacheLine texture_cache_bram[TEXTURE_CACHE_LINES];
    #pragma HLS ARRAY_PARTITION variable=texture_cache_bram complete dim=1
    // #pragma HLS ARRAY_PARTITION variable=texture_cache_bram.data complete dim=1 // 这一行是错误的，请删除

    // 最终输出信号流
    #pragma HLS INTERFACE axis port=tile_processed_signal

    oppc_internal::initialize_tile_buffers(z_buffer_bram, frame_buffer_bram, TILE_PIXELS);
    oppc_internal::initialize_texture_cache(texture_cache_bram, TEXTURE_CACHE_LINES);

    #pragma HLS DATAFLOW

    hls::stream<PixelFragment> valid_fragment_stream("valid_fragment_stream");
    #pragma HLS STREAM variable=valid_fragment_stream depth=32 // 增加流深度，以更好地缓冲

    oppc_internal::read_fragments_producer(fragment_stream, valid_fragment_stream);
    oppc_internal::process_pixels_consumer(
        valid_fragment_stream, config, texture_axi, frame_buffer_bram, z_buffer_bram, texture_cache_bram
    );

    tile_processed_signal.write(true);
}
